import os
import zarr
import numpy as np
from tqdm import tqdm
import cv2
import pickle
#from diffusion_policy.codecs.imagecodecs_numcodecs import register_codecs

class ZarrFun():
    def __init__(self):
        print("ZarrFun Init")
    
    #读取pkl文件并解析
    def read_pkl(self,pkleFile):
        with open(pkleFile, 'rb') as file:
            data = pickle.load(file)
            # imgList = data["camera0_rgb"]
            # for i in range(len(imgList)):
            #     img = imgList[i]
            #     cv2.namedWindow("1",0)
            #     cv2.imshow("1",img)
            #     cv2.waitKey(10)
            #print("成功读取数据长度：", len(data))
            return data
        
    #可视化pkl文件中的图像序列
    def show_pkl_img(self,pkleFile):
        data = self.read_pkl(pkleFile)
        imgList = data["camera0_rgb"]
        for i in range(len(imgList)):
            img = imgList[i]
            cv2.namedWindow("1",0)
            cv2.imshow("1",img)
            cv2.waitKey(10)
        

    def create_empty_zarr(self,storage=None, root=None):
        if root is None:
            if storage is None:
                storage = zarr.MemoryStore()
            root = zarr.group(store=storage)
        data = root.require_group('data', overwrite=False)
        meta = root.require_group('meta', overwrite=False)
        meta.zeros('episode_ends', shape=(0,), dtype=np.int64, compressor=None, overwrite=False)
        data.create_dataset("camera0_rgb", shape=(0, 224, 224, 3), dtype=np.uint8, chunks=(1, 224, 224, 3))
        data.create_dataset("robot0_demo_start_pose", shape=(0, 6), dtype=np.float64, chunks=(1,  6))
        data.create_dataset("robot0_demo_end_pose", shape=(0, 6), dtype=np.float64, chunks=(1,  6))
        data.create_dataset("robot0_eef_pos", shape=(0, 3), dtype=np.float32, chunks=(1,  3))
        data.create_dataset("robot0_eef_rot_axis_angle", shape=(0, 3), dtype=np.float32, chunks=(1,  3))
        data.create_dataset("robot0_gripper_width", shape=(0, 1), dtype=np.float32, chunks=(1,  1))
        return root
    
    def addData(self,tgtZarr,key,value):
        shape = tgtZarr.data[key].shape
        new_shape = (shape[0] + 1,) + shape[1:]
        tgtZarr.data[key].resize(new_shape)
        tgtZarr.data[key][-1] = value


    def zarr_to_zarr(self):
        target_path = '3.zarr.zip'
        tgtZarr = self.create_empty_zarr()

        with zarr.ZipStore(f'cup_in_the_wild.zarr.zip', mode='r') as zip_store:
            print(zip_store.listdir())
            src_root = zarr.group(zip_store)
            episodeEnds = src_root['meta']['episode_ends']
            episode_starts = [0] + episodeEnds[:-1].tolist()
            episode_ends = episodeEnds[:].tolist()
            #print("episode_ends:",episode_ends)
            # 访问每个 episode 的数据
            for i in range(len(episode_ends)):
                if i == 3:
                    break
                start_idx = episode_starts[i]
                end_idx = episode_ends[i]
                cnt = 0
                for j in range(start_idx,end_idx):
                    self.addData(tgtZarr,"robot0_demo_start_pose",src_root['data']['robot0_demo_start_pose'][j])
                    self.addData(tgtZarr,"robot0_demo_end_pose",src_root['data']['robot0_demo_end_pose'][j])
                    self.addData(tgtZarr,"robot0_eef_pos",src_root['data']['robot0_eef_pos'][j])
                    self.addData(tgtZarr,"robot0_eef_rot_axis_angle",src_root['data']['robot0_eef_rot_axis_angle'][j])
                    self.addData(tgtZarr,"robot0_gripper_width",src_root['data']['robot0_gripper_width'][j])
                    self.addData(tgtZarr, "camera0_rgb", src_root['data']['camera0_rgb'][j])
                    cnt = cnt + 1
                # 写espiod
                print("cnt =",cnt)
                shape = tgtZarr.meta["episode_ends"].shape
                new_shape = (shape[0] + 1,) + shape[1:]
                oldCnt = 0
                if shape[0] > 0:
                    oldCnt = tgtZarr.meta["episode_ends"][-1]
                tgtZarr.meta["episode_ends"].resize(new_shape)
                tgtZarr.meta["episode_ends"][-1] = oldCnt + cnt

        print("episode_ends:",tgtZarr.meta["episode_ends"][-1])

        zip_store = zarr.ZipStore(target_path, mode='w')
        zarr.copy_store(tgtZarr.store, zip_store)
        zip_store.close()

    #遍历文件夹，获取pkl文件
    def getFileList(self,folder_path):
        fileList = []
        for root, dirs, files in os.walk(folder_path):
            for file in files:
                if file.endswith('.pkl'):
                    # 获取文件的完整路径
                    file_path = os.path.join(root, file)
                    fileList.append(file_path)
                    #print(file_path)
        return fileList

    def pkl_to_zarr(self,folder_path):
        target_path = 'jt.zarr.zip'
        tgtZarr = self.create_empty_zarr()
        pklList = self.getFileList(folder_path)
        for i in range(len(pklList)):
            print("deal file pklFile: ",pklList[i])
            cnt = 0
            pkl_data = self.read_pkl(pklList[i])

            rgbLen = len(pkl_data["camera0_rgb"])
            posLen = len(pkl_data["robot0_eef_pos"])
            #print(f"###episode:{i}  rgbLen:{rgbLen}  posLen:{posLen}")

            for j in range(len(pkl_data["robot0_eef_pos"])):
                startPos = pkl_data["robot0_eef_pos"][0]
                endPos = pkl_data["robot0_eef_pos"][-1]
                curPos = pkl_data["robot0_eef_pos"][j]
                curImg = pkl_data["camera0_rgb"][j]

                self.addData(tgtZarr,"robot0_demo_start_pose",[num / 1000 for num in startPos]+[3.14, 0.0 ,3.14])
                self.addData(tgtZarr,"robot0_demo_end_pose",[num / 1000 for num in endPos]+[3.14, 0.0 ,3.14])
                self.addData(tgtZarr,"robot0_eef_pos",[num / 1000 for num in curPos])
                self.addData(tgtZarr,"robot0_eef_rot_axis_angle",[3.14, 0.0 ,3.14])
                self.addData(tgtZarr,"robot0_gripper_width",0)
                self.addData(tgtZarr, "camera0_rgb", curImg)
                cnt = cnt + 1

            # 写espiod
            #print("cnt =",cnt)
            shape = tgtZarr.meta["episode_ends"].shape
            new_shape = (shape[0] + 1,) + shape[1:]
            oldCnt = 0
            if shape[0] > 0:
                oldCnt = tgtZarr.meta["episode_ends"][-1]
            tgtZarr.meta["episode_ends"].resize(new_shape)
            tgtZarr.meta["episode_ends"][-1] = oldCnt + cnt

        print("episode_ends:",tgtZarr.meta["episode_ends"][-1])    
        zip_store = zarr.ZipStore(target_path, mode='w')
        zarr.copy_store(tgtZarr.store, zip_store)
        zip_store.close()

zarrFun = ZarrFun()
zarrFun.pkl_to_zarr("./data/acqData/")
#zarrFun.read_pkl("./data/acqData/demo30.pkl")
#zarrFun.getFileList("./data/acqData/")
#zarrFun.show_pkl_img("./data/acqData/demo6.pkl")